﻿using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.DependencyInjection;
using System.Globalization;
using System.Threading.RateLimiting;

namespace Framework.AspNetCore.Extensons
{
	/// <summary>
	/// 速率限制 扩展类
	/// </summary>
	public static class RateLimiterServiceCollectionExtensions
	{
		/// <summary>
		/// 添加 速率限制
		/// </summary>
		/// <param name="services"></param>
		/// <returns></returns>
		public static IServiceCollection AddRateLimiter(this IServiceCollection services)
		{
			//每60秒限制100个请求，滑块添加，分6段
			services.AddRateLimiter(_ =>
			{
				_.RejectionStatusCode = StatusCodes.Status429TooManyRequests;
				_.OnRejected = (context, _) =>
				{
					if (context.Lease.TryGetMetadata(MetadataName.RetryAfter, out var retryAfter))
					{
						context.HttpContext.Response.Headers.RetryAfter =
							((int)retryAfter.TotalSeconds).ToString(NumberFormatInfo.InvariantInfo);
					}
					context.HttpContext.Response.StatusCode = StatusCodes.Status429TooManyRequests;
					context.HttpContext.Response.WriteAsync("请求太多,请稍后再试。");

					return new ValueTask();
				};

				// 全局限流器，全局使用，链式表达式
				_.GlobalLimiter = PartitionedRateLimiter.CreateChained(
				   PartitionedRateLimiter.Create<HttpContext, string>(httpContext =>
				   {
					   var userAgent = httpContext.Request.Headers.UserAgent.ToString();

					   // 滑动窗口限流策略
					   return RateLimitPartition.GetSlidingWindowLimiter
					   (userAgent, _ =>
						   new SlidingWindowRateLimiterOptions
						   {
							   PermitLimit = 100,                           // 每窗口允许最多请求数
							   Window = TimeSpan.FromSeconds(60),     // 每窗口的时间，单位秒
							   SegmentsPerWindow = 6,                       // 滑块添加，分6段
							   QueueProcessingOrder = QueueProcessingOrder.OldestFirst  // 排队请求的处理顺序,优先处理先来的请求
						   });
				   }));
			});
			return services;
		}
	}
}
